"""
股票数据处理模块

该模块提供股票数据的数据库连接、处理和计算功能，主要包括：
1. 数据库连接管理
2. 股票数据表结构管理
3. 自定义字段计算和处理

使用示例：
    from data_processor import process_stock_data
    process_stock_data()
"""
import pandas as pd
from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError
from config import DB_CONFIG

def get_db_engine():
    """
    创建并返回数据库引擎实例
    
    返回:
        Engine: SQLAlchemy数据库引擎实例，配置了连接池和自动重连
    
    配置参数:
        - pool_size: 连接池大小
        - max_overflow: 最大溢出连接数
        - pool_recycle: 连接回收时间(秒)
        - pool_pre_ping: 启用连接前ping检查
    """
    return create_engine(
        f"mysql+mysqlconnector://{DB_CONFIG['user']}:{DB_CONFIG['password']}@"
        f"{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}",
        pool_size=5,
        max_overflow=10,
        pool_recycle=3600,
        pool_pre_ping=True
    )


def calculate_open_close_diff(df):
    """
    计算股票开盘价与收盘价的差值
    
    参数:
        df (DataFrame): 包含open和close字段的股票数据DataFrame
    
    返回:
        DataFrame: 处理后的DataFrame，新增result1字段存储差值
    
    示例:
        >>> df = pd.DataFrame({'open':[10], 'close':[12]})
        >>> calculate_open_close_diff(df)
           open  close  result1
        0    10     12        2
    """
    df = df.copy()
    df['result1'] = df['close'] - df['open']
    return df

def process_stock_data_with_custom_field(field_name, calc_function):
    """
    通用股票数据处理函数，支持自定义字段和计算逻辑
    
    主要功能:
        1. 检查并创建表结构(如不存在)
        2. 动态添加字段(如不存在)
        3. 分批次处理数据
        4. 批量更新数据库
        5. 创建索引优化查询
    
    参数:
        field_name (str): 要添加/更新的字段名称
        calc_function (function): 数据处理函数，接受DataFrame返回处理后的DataFrame
    
    异常:
        SQLAlchemyError: 数据库操作失败时抛出
        Exception: 其他程序异常时抛出
    
    流程说明:
        1. 检查表是否存在，不存在则创建
        2. 检查字段是否存在，不存在则添加
        3. 分批次读取数据(每批100条)
        4. 应用计算函数处理数据
        5. 批量更新数据库
        6. 为字段创建索引
    """
    print('开始处理股票数据...')
    engine = get_db_engine()
    try:
        with engine.connect() as conn:
            # 检查表是否存在
            result = conn.execute(text(
                "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES "
                "WHERE TABLE_NAME = 'stock_index'"
            ))
            if not result.fetchone():
                conn.execute(text(
                    """
                    CREATE TABLE stock_index (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        ts_code VARCHAR(20),
                        trade_date DATE,
                        open FLOAT,
                        high FLOAT,
                        low FLOAT,
                        close FLOAT,
                        pre_close FLOAT,
                        `change` FLOAT,
                        pct_chg FLOAT,
                        vol FLOAT,
                        amount FLOAT
                    )
                    """
                ))
                print('成功创建stock_index表')

            # 动态检查字段是否存在
            result = conn.execute(text(
                "SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS "
                "WHERE TABLE_NAME = 'stock_index' AND COLUMN_NAME = :field"
            ), {'field': field_name})
            
            if not result.fetchone():
                # 动态添加字段
                conn.execute(text(
                    f"ALTER TABLE stock_index "
                    f"ADD COLUMN {field_name} FLOAT NULL AFTER close"
                ))
                print(f'已添加{field_name}字段')

            # 读取数据并使用自定义计算函数
            # 分批次处理数据
            batch_size = 100
            offset = 0
            total_updated = 0
            
            while True:
                # 添加索引查询
                df = pd.read_sql(f'SELECT * FROM stock_index ORDER BY id LIMIT {batch_size} OFFSET {offset}', conn)
                if df.empty:
                    break
                
                # 处理当前批次数据
                processed_df = calc_function(df)
                
                # 使用事务批量更新
                # 批量更新数据
                update_data = [{field_name: row[field_name], 'id': row['id']} for _, row in processed_df.iterrows()]
                conn.execute(text(
                    f"UPDATE stock_index SET {field_name} = :{field_name} WHERE id = :id"
                ), update_data)
                conn.commit()
                print(f'成功提交{len(update_data)}条更新')
                
                total_updated += len(processed_df)
                offset += batch_size
            
            print(f'成功更新{total_updated}条记录')
            
            # 添加索引优化查询
            index_check = conn.execute(text(
                "SHOW INDEX FROM stock_index WHERE Column_name = :field"),
                {'field': field_name}
            )
            if not index_check.fetchone():
                conn.execute(text(f"CREATE INDEX idx_{field_name} ON stock_index({field_name})").execution_options(autocommit=True))
            
            conn.commit()

    except SQLAlchemyError as e:
        print(f'数据库操作失败：{str(e)}')
        conn.rollback()
        print(f'事务回滚：{str(e)}')
    except Exception as e:
        print(f'程序异常：{str(e)}')
        raise

def process_stock_data():
    """
    默认股票数据处理函数
    
    使用calculate_open_close_diff作为默认计算函数，
    处理result1字段(收盘价-开盘价)
    
    这是process_stock_data_with_custom_field的简化封装
    """
    process_stock_data_with_custom_field('result1', calculate_open_close_diff)

if __name__ == '__main__':
    process_stock_data()